#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2024/12/01 22:55
@Author  : thezehui@gmail.com
@File    : chat.py
"""
from typing import List, Tuple

import tiktoken
from langchain_community.chat_models.tongyi import ChatTongyi
from langchain_core.messages import BaseMessage, get_buffer_string

from internal.core.language_model.entities.model_entity import BaseLanguageModel


class Chat(ChatTongyi, BaseLanguageModel):
    """通义千问聊天模型"""
    
    def _get_encoding_model(self) -> Tuple[str, tiktoken.Encoding]:
        """获取编码模型名字+模型函数，使用cl100k_base作为编码器
        
        Returns:
            模型名称和编码器的元组
        """
        # 使用通用的cl100k_base编码器，因为通义千问没有专门的tiktoken编码器
        model = "cl100k_base"
        return model, tiktoken.get_encoding(model)

    def get_num_tokens_from_messages(self, messages: List[BaseMessage]) -> int:
        """获取消息token数量的方法，使用cl100k_base编码器计算token数量
        
        Args:
            messages: 消息列表
            
        Returns:
            消息的token数量
        """
        # 获取编码器
        _, encoding = self._get_encoding_model()

        # 计算所有消息的token数量
        num_tokens = 0
        for message in messages:
            # 获取消息的文本内容
            message_text = get_buffer_string([message])
            # 使用编码器计算token数量
            num_tokens += len(encoding.encode(message_text))

            # 如果消息有工具调用，也需要计算工具调用的token数量
            if hasattr(message, "tool_calls") and message.tool_calls:
                for tool_call in message.tool_calls:
                    # 工具名称
                    if "name" in tool_call:
                        num_tokens += len(encoding.encode(tool_call["name"]))
                    # 工具参数 (转为字符串计算)
                    if "args" in tool_call:
                        args_str = str(tool_call["args"])
                        num_tokens += len(encoding.encode(args_str))

        return num_tokens
